package com.mall.hadoop.format;

import java.io.DataOutputStream;
import java.io.IOException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.io.compress.CompressionCodec;
import org.apache.hadoop.io.compress.GzipCodec;
import org.apache.hadoop.mapreduce.OutputCommitter;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputCommitter;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.util.ReflectionUtils;


/**
 * 自定义抽象类MultipleOutputFormat继承抽象类FileOutputFormat
 * 需要在主程序中实现generateFileNameForKeyValue来获取文件名 
 * @author tgy
 *
 * @param <K>
 * @param <V>
 */
public abstract class MultipleOutputFormat<K extends WritableComparable<?>, V extends Writable> extends FileOutputFormat<K, V> {  
	
	private MultiRecordWriter writer = null;
	public RecordWriter<K, V> getRecordWriter(TaskAttemptContext job) throws IOException, InterruptedException {
		if(writer==null){
			//通过getTaskOutputPath获取output路径
			writer = new MultiRecordWriter(job, getTaskOutputPath(job));
		}
		return writer;
	}
	private Path getTaskOutputPath(TaskAttemptContext job) throws IOException{
		Path workPath = null;  
        OutputCommitter committer = super.getOutputCommitter(job);  
        if (committer instanceof FileOutputCommitter) {  
            workPath = ((FileOutputCommitter) committer).getWorkPath();  
        } else {  
            Path outputPath = super.getOutputPath(job);  
            if (outputPath == null) {  
                throw new IOException("Undefined job output-path");  
            }  
            workPath = outputPath;  
        }  
        return workPath;
	}
	 //通过key, value, conf来确定输出文件名（含扩展名），返回值就是文件名。可以根据key,value来判断  
    protected abstract String generateFileNameForKeyValue(K key, V value, Configuration conf);
    
	public class MultiRecordWriter extends RecordWriter<K, V>{
		//RecordWriter的缓存
		private Map<String, RecordWriter<K, V>> recordWriters = null;
		private TaskAttemptContext job = null;
		//输出目录
		private Path workPath = null;
		public MultiRecordWriter(TaskAttemptContext job,Path workPath) {
			super();
			this.job = job;
			this.workPath = workPath;
			recordWriters = new HashMap<String, RecordWriter<K, V>>();
		}
		
		@Override
		public void write(K key, V value) throws IOException, InterruptedException {
			//得出输出文件名
			String baseName = generateFileNameForKeyValue(key, value, job.getConfiguration());
			RecordWriter<K, V> rw = this.recordWriters.get(baseName);
			if(rw==null){
				rw = getBaseRecordWriter(job, baseName);
				this.recordWriters.put(baseName, rw);
			}
			rw.write(key, value);
		}

		//关闭，可能是多个文件进行关闭，所有采用循环   
        //recordWriters.values() 就是指的getBaseRecordWriter返回的值
		@Override
		public void close(TaskAttemptContext context) throws IOException, InterruptedException {
			Iterator<RecordWriter<K, V>> values = this.recordWriters.values().iterator();
			while(values.hasNext()){
				values.next().close(context);
			}
			this.recordWriters.clear();
		}
		
		private RecordWriter<K, V> getBaseRecordWriter(TaskAttemptContext job,String name) throws IOException, InterruptedException{
			//获取配置文件
			Configuration config = job.getConfiguration();
			//查看是否使用解码器
			boolean flag = getCompressOutput(job);
			String keyValueSeparator = ",";
			RecordWriter<K, V> recordWriter = null; 
			if(flag){
				Class<? extends CompressionCodec> codecClass = getOutputCompressorClass(job, GzipCodec.class);
				CompressionCodec codec = ReflectionUtils.newInstance(codecClass, config);
				Path file = new Path(workPath,name+codec.getDefaultExtension());
				FSDataOutputStream fileOut = file.getFileSystem(config).create(file, false);
				recordWriter = new LineRecordWriter<K, V>(new DataOutputStream(codec.createOutputStream(fileOut)), keyValueSeparator);
			} else {  //如果不使用解码器
				Path file = new Path(workPath, name);
				FSDataOutputStream fileOut = file.getFileSystem(config).create(file, false);
				//这里使用自己的OutputFormat  
				recordWriter = new LineRecordWriter<K, V>(fileOut, keyValueSeparator);
			} 
            return recordWriter;
		}
	}
}
